{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-26T04:12:13.603867Z",
     "start_time": "2021-07-26T04:12:13.591802Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.utils import add_remaining_self_loops\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv,GATConv\n",
    "import numpy as np\n",
    "import pickle\n",
    "from torch_geometric.nn import JumpingKnowledge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-26T04:12:14.576978Z",
     "start_time": "2021-07-26T04:12:14.545449Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(edge_index=[2, 13566], label_feature=[2708, 7], num_label=7, x=[2708, 1433], y=[2708])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open(\"data/cora/cora.pkl\",'rb') as f:\n",
    "    cora = pickle.load(f)\n",
    "cora"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-26T03:55:40.164533Z",
     "start_time": "2021-07-26T03:55:40.149962Z"
    }
   },
   "outputs": [],
   "source": [
    "# JK-Nets（6层）\n",
    "class JKNet(torch.nn.Module):\n",
    "    def __init__(self, dataset, mode='cat', num_layers=6, hidden=16):\n",
    "        super(JKNet, self).__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.mode = mode\n",
    "\n",
    "        self.conv0 = GCNConv(dataset.x.shape[1], hidden)\n",
    "        self.dropout0 = torch.nn.Dropout(p=0.5)\n",
    "\n",
    "        for i in range(1, self.num_layers):\n",
    "            setattr(self, 'conv{}'.format(i), GCNConv(hidden, hidden))\n",
    "            setattr(self, 'dropout{}'.format(i), torch.nn.Dropout(p=0.5))\n",
    "\n",
    "        self.jk = JumpingKnowledge(mode=mode)\n",
    "        if mode == 'max':\n",
    "            self.fc = torch.nn.Linear(hidden, 7)\n",
    "        elif mode == 'cat':\n",
    "            self.fc = torch.nn.Linear(num_layers * hidden, 7)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index = data.x, data.edge_index\n",
    "\n",
    "        layer_out = []  # 保存每一层的结果\n",
    "        for i in range(self.num_layers):\n",
    "            conv = getattr(self, 'conv{}'.format(i))\n",
    "            dropout = getattr(self, 'dropout{}'.format(i))\n",
    "            x = dropout(F.relu(conv(x, edge_index)))\n",
    "            layer_out.append(x)\n",
    "\n",
    "        h = self.jk(layer_out)  # JK层\n",
    "\n",
    "        h = self.fc(h)\n",
    "        h = F.log_softmax(h, dim=1)\n",
    "\n",
    "        return h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-26T03:55:41.719507Z",
     "start_time": "2021-07-26T03:55:41.710606Z"
    }
   },
   "outputs": [],
   "source": [
    "mask = torch.randperm(cora.x.shape[0])\n",
    "\n",
    "# train_mask = mask[:140]\n",
    "# val_mask = mask[140:640]\n",
    "# test_mask = mask[1708:2708]\n",
    "\n",
    "train_mask = mask[:1208]\n",
    "unlabeled_mask = mask[1208:]\n",
    "val_mask = mask[1208:1708]\n",
    "test_mask = mask[1708:2708]\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-26T04:13:39.780799Z",
     "start_time": "2021-07-26T04:13:37.032925Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 loss: 1.9655\n",
      "epoch: 1 loss: 1.9483\n",
      "epoch: 2 loss: 1.9257\n",
      "epoch: 3 loss: 1.8921\n",
      "epoch: 4 loss: 1.8522\n",
      "epoch: 5 loss: 1.8111\n",
      "epoch: 6 loss: 1.7671\n",
      "epoch: 7 loss: 1.7061\n",
      "epoch: 8 loss: 1.6687\n",
      "epoch: 9 loss: 1.6101\n",
      "Accuracy: 0.3700\n",
      "epoch: 10 loss: 1.5679\n",
      "epoch: 11 loss: 1.5115\n",
      "epoch: 12 loss: 1.4605\n",
      "epoch: 13 loss: 1.4227\n",
      "epoch: 14 loss: 1.3382\n",
      "epoch: 15 loss: 1.2972\n",
      "epoch: 16 loss: 1.2323\n",
      "epoch: 17 loss: 1.1907\n",
      "epoch: 18 loss: 1.1465\n",
      "epoch: 19 loss: 1.1031\n",
      "Accuracy: 0.6690\n",
      "epoch: 20 loss: 1.0557\n",
      "epoch: 21 loss: 0.9841\n",
      "epoch: 22 loss: 0.9788\n",
      "epoch: 23 loss: 0.9244\n",
      "epoch: 24 loss: 0.9038\n",
      "epoch: 25 loss: 0.8528\n",
      "epoch: 26 loss: 0.8319\n",
      "epoch: 27 loss: 0.7697\n",
      "epoch: 28 loss: 0.7472\n",
      "epoch: 29 loss: 0.7303\n",
      "Accuracy: 0.7930\n",
      "epoch: 30 loss: 0.6858\n",
      "epoch: 31 loss: 0.7098\n",
      "epoch: 32 loss: 0.6562\n",
      "epoch: 33 loss: 0.6713\n",
      "epoch: 34 loss: 0.6295\n",
      "epoch: 35 loss: 0.5637\n",
      "epoch: 36 loss: 0.5653\n",
      "epoch: 37 loss: 0.5209\n",
      "epoch: 38 loss: 0.5217\n",
      "epoch: 39 loss: 0.5168\n",
      "Accuracy: 0.8460\n",
      "epoch: 40 loss: 0.5146\n",
      "epoch: 41 loss: 0.4886\n",
      "epoch: 42 loss: 0.4428\n",
      "epoch: 43 loss: 0.4477\n",
      "epoch: 44 loss: 0.4384\n",
      "epoch: 45 loss: 0.4579\n",
      "epoch: 46 loss: 0.4472\n",
      "epoch: 47 loss: 0.4064\n",
      "epoch: 48 loss: 0.3973\n",
      "epoch: 49 loss: 0.3785\n",
      "Accuracy: 0.8660\n",
      "epoch: 50 loss: 0.3795\n",
      "epoch: 51 loss: 0.3985\n",
      "epoch: 52 loss: 0.3436\n",
      "epoch: 53 loss: 0.3861\n",
      "epoch: 54 loss: 0.3701\n",
      "epoch: 55 loss: 0.3498\n",
      "epoch: 56 loss: 0.3410\n",
      "epoch: 57 loss: 0.3206\n",
      "epoch: 58 loss: 0.3317\n",
      "epoch: 59 loss: 0.3176\n",
      "Accuracy: 0.8710\n",
      "epoch: 60 loss: 0.3492\n",
      "epoch: 61 loss: 0.3248\n",
      "epoch: 62 loss: 0.3038\n",
      "epoch: 63 loss: 0.2645\n",
      "epoch: 64 loss: 0.2875\n",
      "epoch: 65 loss: 0.3229\n",
      "epoch: 66 loss: 0.3034\n",
      "epoch: 67 loss: 0.2792\n",
      "epoch: 68 loss: 0.2803\n",
      "epoch: 69 loss: 0.2716\n",
      "Accuracy: 0.8640\n",
      "epoch: 70 loss: 0.2718\n",
      "epoch: 71 loss: 0.2508\n",
      "epoch: 72 loss: 0.2549\n",
      "epoch: 73 loss: 0.2450\n",
      "epoch: 74 loss: 0.2497\n",
      "epoch: 75 loss: 0.2774\n",
      "epoch: 76 loss: 0.2655\n",
      "epoch: 77 loss: 0.2682\n",
      "epoch: 78 loss: 0.2692\n",
      "epoch: 79 loss: 0.2571\n",
      "Accuracy: 0.8630\n",
      "epoch: 80 loss: 0.2600\n",
      "epoch: 81 loss: 0.2516\n",
      "epoch: 82 loss: 0.2325\n",
      "epoch: 83 loss: 0.2530\n",
      "epoch: 84 loss: 0.2286\n",
      "epoch: 85 loss: 0.2427\n",
      "epoch: 86 loss: 0.2228\n",
      "epoch: 87 loss: 0.2302\n",
      "epoch: 88 loss: 0.2592\n",
      "epoch: 89 loss: 0.2417\n",
      "Accuracy: 0.8590\n",
      "epoch: 90 loss: 0.2142\n",
      "epoch: 91 loss: 0.2291\n",
      "epoch: 92 loss: 0.2217\n",
      "epoch: 93 loss: 0.2107\n",
      "epoch: 94 loss: 0.2106\n",
      "epoch: 95 loss: 0.2114\n",
      "epoch: 96 loss: 0.2187\n",
      "epoch: 97 loss: 0.2188\n",
      "epoch: 98 loss: 0.2277\n",
      "epoch: 99 loss: 0.2093\n",
      "Accuracy: 0.8600\n",
      "epoch: 100 loss: 0.2100\n",
      "epoch: 101 loss: 0.1940\n",
      "epoch: 102 loss: 0.2169\n",
      "epoch: 103 loss: 0.1992\n",
      "epoch: 104 loss: 0.1940\n",
      "epoch: 105 loss: 0.2146\n",
      "epoch: 106 loss: 0.1762\n",
      "epoch: 107 loss: 0.1983\n",
      "epoch: 108 loss: 0.1892\n",
      "epoch: 109 loss: 0.1914\n",
      "Accuracy: 0.8610\n",
      "epoch: 110 loss: 0.1859\n",
      "epoch: 111 loss: 0.1937\n",
      "epoch: 112 loss: 0.2023\n",
      "epoch: 113 loss: 0.1849\n",
      "epoch: 114 loss: 0.1914\n",
      "epoch: 115 loss: 0.1820\n",
      "epoch: 116 loss: 0.1674\n",
      "epoch: 117 loss: 0.1839\n",
      "epoch: 118 loss: 0.1845\n",
      "epoch: 119 loss: 0.1757\n",
      "Accuracy: 0.8570\n",
      "epoch: 120 loss: 0.1868\n",
      "epoch: 121 loss: 0.1834\n",
      "epoch: 122 loss: 0.1762\n",
      "epoch: 123 loss: 0.1663\n",
      "epoch: 124 loss: 0.1806\n",
      "epoch: 125 loss: 0.1724\n",
      "epoch: 126 loss: 0.1854\n",
      "epoch: 127 loss: 0.1770\n",
      "epoch: 128 loss: 0.1668\n",
      "epoch: 129 loss: 0.1582\n",
      "Accuracy: 0.8590\n",
      "epoch: 130 loss: 0.1641\n",
      "epoch: 131 loss: 0.1710\n",
      "epoch: 132 loss: 0.1742\n",
      "epoch: 133 loss: 0.1578\n",
      "epoch: 134 loss: 0.1657\n",
      "epoch: 135 loss: 0.1660\n",
      "epoch: 136 loss: 0.1528\n",
      "epoch: 137 loss: 0.1616\n",
      "epoch: 138 loss: 0.1458\n",
      "epoch: 139 loss: 0.1878\n",
      "Accuracy: 0.8630\n",
      "epoch: 140 loss: 0.1525\n",
      "epoch: 141 loss: 0.1508\n",
      "epoch: 142 loss: 0.1477\n",
      "epoch: 143 loss: 0.1636\n",
      "epoch: 144 loss: 0.1558\n",
      "epoch: 145 loss: 0.1649\n",
      "epoch: 146 loss: 0.1412\n",
      "epoch: 147 loss: 0.1443\n",
      "epoch: 148 loss: 0.1549\n",
      "epoch: 149 loss: 0.1492\n",
      "Accuracy: 0.8590\n",
      "epoch: 150 loss: 0.1677\n",
      "epoch: 151 loss: 0.1573\n",
      "epoch: 152 loss: 0.1597\n",
      "epoch: 153 loss: 0.1653\n",
      "epoch: 154 loss: 0.1701\n",
      "epoch: 155 loss: 0.1444\n",
      "epoch: 156 loss: 0.1503\n",
      "epoch: 157 loss: 0.1343\n",
      "epoch: 158 loss: 0.1493\n",
      "epoch: 159 loss: 0.1396\n",
      "Accuracy: 0.8580\n",
      "epoch: 160 loss: 0.1509\n",
      "epoch: 161 loss: 0.1563\n",
      "epoch: 162 loss: 0.1591\n",
      "epoch: 163 loss: 0.1438\n",
      "epoch: 164 loss: 0.1388\n",
      "epoch: 165 loss: 0.1471\n",
      "epoch: 166 loss: 0.1520\n",
      "epoch: 167 loss: 0.1667\n",
      "epoch: 168 loss: 0.1475\n",
      "epoch: 169 loss: 0.1502\n",
      "Accuracy: 0.8560\n",
      "epoch: 170 loss: 0.1564\n",
      "epoch: 171 loss: 0.1385\n",
      "epoch: 172 loss: 0.1325\n",
      "epoch: 173 loss: 0.1407\n",
      "epoch: 174 loss: 0.1456\n",
      "epoch: 175 loss: 0.1391\n",
      "epoch: 176 loss: 0.1278\n",
      "epoch: 177 loss: 0.1225\n",
      "epoch: 178 loss: 0.1438\n",
      "epoch: 179 loss: 0.1568\n",
      "Accuracy: 0.8620\n",
      "epoch: 180 loss: 0.1250\n",
      "epoch: 181 loss: 0.1443\n",
      "epoch: 182 loss: 0.1267\n",
      "epoch: 183 loss: 0.1372\n",
      "epoch: 184 loss: 0.1327\n",
      "epoch: 185 loss: 0.1410\n",
      "epoch: 186 loss: 0.1322\n",
      "epoch: 187 loss: 0.1387\n",
      "epoch: 188 loss: 0.1366\n",
      "epoch: 189 loss: 0.1341\n",
      "Accuracy: 0.8530\n",
      "epoch: 190 loss: 0.1385\n",
      "epoch: 191 loss: 0.1081\n",
      "epoch: 192 loss: 0.1274\n",
      "epoch: 193 loss: 0.1347\n",
      "epoch: 194 loss: 0.1318\n",
      "epoch: 195 loss: 0.1158\n",
      "epoch: 196 loss: 0.1312\n",
      "epoch: 197 loss: 0.1521\n",
      "epoch: 198 loss: 0.1417\n",
      "epoch: 199 loss: 0.1367\n",
      "Accuracy: 0.8560\n"
     ]
    }
   ],
   "source": [
    "model = JKNet(cora).to(device)\n",
    "cora = cora.to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=5e-3, weight_decay=5e-4)\n",
    "\n",
    "for epoch in range(200):\n",
    "    optimizer.zero_grad()\n",
    "    out = model(cora)\n",
    "    loss = F.nll_loss(out[train_mask], cora.y[train_mask])\n",
    "    print('epoch: %d loss: %.4f' %(epoch, loss))\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if((epoch + 1)% 10 == 0):\n",
    "        model.eval()\n",
    "        _, pred = model(cora).max(dim=1)\n",
    "        correct = int(pred[test_mask].eq(cora.y[test_mask]).sum().item())\n",
    "        acc = correct / len(test_mask)\n",
    "        print('Accuracy: {:.4f}'.format(acc))\n",
    "        model.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "song",
   "language": "python",
   "name": "song"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
